import pandas as pd
import requests
from datetime import datetime
import os


# Step 1: Generate trading days
def generate_trading_days(year):
    start_date = f"{year}-01-01"
    end_date = f"{year}-12-31"
    date_range = pd.date_range(start=start_date, end=end_date)

    # Define holidays for the given year
    holidays = [
        f"{year}-01-01",  # New Year's Day
        f"{year}-01-28", f"{year}-01-29", f"{year}-01-30", f"{year}-01-31", f"{year}-02-01", f"{year}-02-02",
        f"{year}-02-03", f"{year}-02-04",  # Spring Festival
        f"{year}-04-04", f"{year}-04-05", f"{year}-04-06",  # Qingming Festival
        f"{year}-05-01", f"{year}-05-02", f"{year}-05-03", f"{year}-05-04", f"{year}-05-05",  # Labor Day
        f"{year}-06-01", f"{year}-06-02", f"{year}-06-03",  # Dragon Boat Festival
        f"{year}-10-01", f"{year}-10-02", f"{year}-10-03", f"{year}-10-04", f"{year}-10-05", f"{year}-10-06",
        f"{year}-10-07", f"{year}-10-08",  # National Day and Mid-Autumn Festival
    ]

    # Filter trading days
    trading_days = []
    for date in date_range:
        if date.strftime('%Y-%m-%d') not in holidays and date.weekday() < 5:
            trading_days.append(date.strftime('%Y-%m-%d'))
    return trading_days


# Step 2: Download files for each trading day
def download_market_data(trading_days, base_url_template, save_directory):
    if not os.path.exists(save_directory):
        os.makedirs(save_directory)

    current_date_str = datetime.now().date().strftime('%Y-%m-%d')
    if current_date_str in trading_days:
        try:
            # Parse date components
            year, month, day = current_date_str.split('-')

            # Generate URL
            url = base_url_template.format(year_month=f"{year}{month}", day=day)

            # Generate filename
            filename = os.path.join(save_directory, f"{year[-2:]}{month}{day}IM持仓.csv")

            # Download file
            response = requests.get(url, stream=True)
            response.raise_for_status()  # Check for HTTP errors

            # Save file
            with open(filename, 'wb') as file:
                for chunk in response.iter_content(chunk_size=8192):
                    file.write(chunk)
            print(f"Downloaded: {filename}")

        except requests.exceptions.RequestException as e:
            print(f"Failed to download {url}: {e}")
        except Exception as e:
            print(f"An error occurred: {e}")


if __name__ == "__main__":
    # Configuration
    YEAR = 2025  # Target year
    BASE_URL_TEMPLATE = "http://www.cffex.com.cn/sj/ccpm/{year_month}/{day}/IM_1.csv"
    SAVE_DIRECTORY = "downloaded_data"

    # Generate trading days
    trading_days = generate_trading_days(YEAR)
    print(f"Generated {len(trading_days)} trading days for {YEAR}.")

    # Download market data
    download_market_data(trading_days, BASE_URL_TEMPLATE, SAVE_DIRECTORY)